library(sf)
library(gstat)
library(future)
library(furrr)
library(readr)
library(dplyr)

# ✅ Setup
options(future.globals.maxSize = 10 * 1024^3)  
plan(multisession, workers = 2)

# ✅ Load grid (with gid)
grid <- st_read("/home/fbai9728/tmp-2/priogrid_mex/priogrid_mex.shp")

# ✅ Find all monthly data files
month_files <- list.files("/home/fbai9728/tmp-2/monthly_chunks/", pattern = "\\.rds$", full.names = TRUE)

# ✅ IDW interpolation function
idw_interpolation <- function(data_subset, group_label) {
  message(paste("Processing:", group_label))
  
  if (nrow(data_subset) < 3) {
    message(paste("⚠️ Skipping", group_label, "- insufficient data"))
    return(NULL)
  }
  
  this_points <- st_as_sf(data_subset, coords = c("lon", "lat"), crs = 4326)
  this_idw <- idw(scalar_sum ~ 1, this_points, newdata = grid, idp = 2)
  
  # Convert result to dataframe
  this_idw_df <- as.data.frame(this_idw)
  this_idw_df$gid <- grid$gid
  this_idw_df$group <- group_label
  this_idw_df <- this_idw_df %>% select(gid, group, var1.pred)
  
  return(this_idw_df)
}

# ✅ Paths for output files
daily_output_path <- "/home/fbai9728/tmp-2/idw_daily_"
monthly_output_path <- "/home/fbai9728/tmp-2/idw_monthly_"

# ✅ Process each month
for (month_file in month_files) {
  month_data <- readRDS(month_file)
  month_label <- gsub("\\.rds$", "", basename(month_file))  # Extract YYYY-MM

  # ✅ Check if monthly file exists - SKIP if done
  monthly_file <- paste0(monthly_output_path, month_label, ".csv")
  if (file.exists(monthly_file)) {
    message(paste("✅ Monthly IDW already exists, skipping:", month_label))
  } else {
    # ✅ Monthly IDW processing (full month)
    monthly_idw_df <- idw_interpolation(month_data, month_label)
    if (!is.null(monthly_idw_df)) {
      write_csv(monthly_idw_df, monthly_file)
      message(paste("✅ Monthly IDW saved for", month_label))
    } else {
      message(paste("⚠️ Monthly IDW skipped for", month_label))
    }
  }

  # ✅ Split data by day for daily interpolation
  daily_data <- split(month_data, month_data$date)
  
  # ✅ Process each day within the month
  daily_results <- list()
  for (day in names(daily_data)) {
    daily_file <- paste0(daily_output_path, day, ".csv")
    
    # ✅ Check if daily file exists - SKIP if done
    if (file.exists(daily_file)) {
      message(paste("✅ Daily IDW already exists, skipping:", day))
      next
    }
    
    # ✅ Run IDW for that day
    daily_idw <- idw_interpolation(daily_data[[day]], day)
    if (!is.null(daily_idw)) {
      write_csv(daily_idw, daily_file)
      message(paste("✅ Daily IDW saved for", day))
    } else {
      message(paste("⚠️ Daily IDW skipped for", day))
    }
  }
}

print("✅ Finished all daily and monthly IDW runs.")
